#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/12/29 22:33
# @Author  : 新猿异码
# @Site    : https://gitee.com/chshao/autotest
# @CSDN    : https://blog.csdn.net/m0_37576542?type=blog
# @File    : minicap.py
# @Desc    : 模块功能描述
# @Software: PyCharm
import os
import re
import time
import struct
import socket
import datetime
from autotest.buildin.PopenOS import Popen
from autotest.kernel.logger.AutoLogging import Log

from autotest.module.android.adb import AdbClient
from autotest.module.android.decorator import asynch
from autotest.module.android.decorator import singleton
from autotest.module.android.decorator import except_wrapper

STF_LIB_PATH = os.path.join(os.path.dirname(__file__), "stf_libs")
PORT1717 = 1717  # adb端口转发，将minicap服务与此端口进行转发


class Banner(object):
    def __init__(self):
        self.__banner = dict(zip(self.keys, [0] * len(self.keys)))

    def __setitem__(self, key, value):
        self.__banner[key] = value

    def __getitem__(self, key):
        return self.__banner[key]

    @property
    def keys(self):
        # 一旦与minicap建立连接，minicap首先会推送长度为24字节的global header，header包含以下9个内容
        return ['version', 'length', 'pid', 'realWidth', 'realHeight', 'virtualWidth', 'virtualHeight', 'orientation',
                'quirks']

    def __str__(self):
        return str(self.__banner)


@singleton
class Minicap(object):
    DEVICE_DIR = "/data/local/tmp"  # minicap在设备上的存放路径
    COMMAND = f"LD_LIBRARY_PATH={DEVICE_DIR} {DEVICE_DIR}/minicap"
    VERSION = 5  # minicap的版本

    def __init__(self, adb: AdbClient = None, sn: str = None, save_path=None, callback=None):
        """
        :param adb: AdbClient实例对象
        :param sn: 设备SN号
        :param save_path: 截图保存路径
        :param callback: 收到一张完整图片数据后的回调函数
        """
        self.adb = adb or AdbClient()
        self.save_path = save_path
        self.callback = callback
        self.sn = sn
        if not sn:
            self.adb.must_one_device(sn)
            self.sn = self.adb.device_list()[0]
        if self.save_path is not None:
            if not os.path.exists(self.save_path):
                os.makedirs(self.save_path)

        # 接收图片数据流的标志
        self.capture_stream_flag = False
        # 调用take_screenshot函数时，截图保存在安卓设备的路径
        self.adb.makedirs(self.sn, f"{self.DEVICE_DIR}/minicap-images")
        # 获取设备屏幕尺寸
        self.display = self.adb.get_display(self.sn)

    def __str__(self):
        return f"{self.__class__.__name__}(SN={self.sn})."

    def install(self):
        """
        安装minicap程序，minicap源码获取地址：https://github.com/openstf/
        :return:
        """
        # 先检查是否已经存在，且是最新版本，如果是则不需要安装
        if self.adb.exists_file(self.sn, f"{self.DEVICE_DIR}/minicap") and \
                self.adb.exists_file(self.sn, f"{self.DEVICE_DIR}/minicap.so"):
            msg = self.adb.shell(self.sn, f"{self.COMMAND} -v 2>&1")
            m = re.match("version:([0-9])", msg)
            version = int(m.group(1)) if m else -1
            if version >= self.VERSION:
                Log.info(f"SN={self.sn}, minicap installed already")
                return True
            self.adb.shell(self.sn, f"rm -rf {self.DEVICE_DIR}/minicap")
            self.adb.shell(self.sn, f"rm -rf {self.DEVICE_DIR}/minicap.so")
        # 没有安装或为旧版本，进行安装
        abi = self.adb.shell(self.sn, "getprop ro.product.cpu.abi").strip("\n")
        pre = self.adb.shell(self.sn, "getprop ro.build.version.preview_sdk").strip("\n")
        rel = self.adb.shell(self.sn, "getprop ro.build.version.release").strip("\n")
        sdk = int(self.adb.shell(self.sn, "getprop ro.build.version.sdk").strip("\n"))
        if pre.isdigit():
            if int(pre) > 0:
                sdk += 1
        # 根据不同的Android SDK版本选择对应的minicap程序
        bin_file = "minicap" if sdk >= 16 else "minicap-nopie"
        # 将minicap推送至手机目标路径下，并赋予读、写和执行权限
        local = os.path.join(STF_LIB_PATH, abi, bin_file)
        self.adb.push(local, f"{self.DEVICE_DIR}/minicap", self.sn)
        self.adb.shell(self.sn, f"chmod 755 {self.DEVICE_DIR}/minicap")

        # 根据不同的Android SDK版本和CPU类型选择对应的minicap.so程序
        pattern = os.path.join(STF_LIB_PATH, 'minicap-shared/aosp/libs/android-%s/%s/minicap.so')
        so_file_path = pattern % (sdk, abi)
        if not os.path.isfile(so_file_path):
            so_file_path = pattern % (rel, abi)
        # 将minicap.so推送至手机目标路径下，并赋予读、写和执行权限
        self.adb.push(so_file_path, f"{self.DEVICE_DIR}/minicap.so", self.sn)
        self.adb.shell(self.sn, f"chmod 755 {self.DEVICE_DIR}/minicap.so")
        Log.info(f"SN={self.sn}, minicap install finished")
        return True

    @except_wrapper
    def start_server(self):
        """
        启动安卓设备侧的minicap服务
        :return:
        """
        # 先查看是否已启动，如果是则先关闭，防止发生过横竖屏切换导致的截图不正确
        pid = self.get_pid()
        if pid is not None:
            self.adb.shell(self.sn, f"kill {pid}")
        # 启动minicap服务
        display = "x".join(self.display)
        proc = Popen(f"adb -s {self.sn} shell {self.COMMAND} -P {display}@{display}/0", daemonic=True)
        proc.wait_until("Server start", timeout=2)
        pid = self.get_pid()
        if pid is None:
            Log.error(f"SN={self.sn}, minicap start failed!")
            return False
        Log.info(f"SN={self.sn}, minicap start success!")
        Popen(f"adb -s {self.sn} forward tcp:{PORT1717} localabstract:minicap")
        return True

    @except_wrapper
    def stop_server(self):
        """
        关闭安卓设备侧的minicap服务
        :return:
        """
        pid = self.get_pid()
        if pid is not None:
            self.adb.shell(self.sn, f"kill {pid}")
            Log.info(f"SN={self.sn}, process of minicap stopped!")
        return True

    def get_pid(self):
        """
        获取minicap进程号
        :return:
        """
        p = os.popen(f"adb -s {self.sn} shell ps | findstr minicap")
        pid_re = re.search("[0-9]+", p.read())
        if pid_re is None:
            # Log.info(f"SN={self.sn}, has not found minicap process")
            return None
        pid = pid_re.group()
        return int(pid)

    @except_wrapper
    def take_screenshot(self, save_path=None, filename=None, mouse=False, is_delete=False):
        """
        截屏
        :param save_path: 本地保存路径，截图会被从设备侧导出到电脑本地，None表示不导出
        :param filename: 图片名称
        :param mouse: 是否给图片添加鼠标图标，方便识别鼠标在哪里做了点击事件，暂未实现
        :param is_delete: 是否删除设备侧的截图
        :return:
        """
        display = "x".join(self.adb.get_display(self.sn))
        filename = datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f') + ".png" if filename is None else filename
        png_file = f"{self.DEVICE_DIR}/minicap-images/{filename}"
        save_file = os.path.join(save_path, filename)
        # 使用minicap截图
        self.adb.shell(self.sn, f"{self.COMMAND} -P {display}@{display}/0 -s > {png_file}")
        # 导出至电脑本地路径
        if save_path is not None:
            self.adb.pull(png_file, save_file)
        if is_delete:
            self.adb.shell(self.sn, f"rm -rf {png_file}")
        return save_file

    @asynch()
    def register_orientation_watcher(self):
        """
        监听设备横竖屏切换事件，若发生横竖屏切换，则重新启动minicap服务
        :return:
        """
        while self.capture_stream_flag:
            display = self.adb.get_display(self.sn)
            # 不相等说明发生了横竖屏切换
            if display[0] != self.display[0]:
                self.display = display
                Log.info(f"SN={self.sn}, 发生了横竖屏切换")
                self.stop_server()
                self.start_server()
                self.start_capture_stream()
            time.sleep(0.2)

    def stop_capture_stream(self):
        """
        停止从minicap服务获取截图数据流
        :return:
        """
        self.capture_stream_flag = False

    def start_capture_stream(self):
        """
        跟minicap服务建立连接，并获取截图数据流
        :return:
        """
        self.capture_stream_flag = True
        self.register_orientation_watcher()
        client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client.settimeout(3)
        if client.connect_ex(("127.0.0.1", PORT1717)) != 0:
            Log.error(f"SN={self.sn}, connect minicap fail.")
            raise ConnectionError(f"SN={self.sn}, connect minicap fail.")
        # 开始接收数据流，一旦与minicap服务建立连接，minicap首先会推送长度为24字节的global header，
        # global header只会推送一次，后续推送的数据不会再包括global header，而是不断的推送实时图片流数据，直到客户端关闭socket连接
        self.__get_stream(client)

    @asynch()
    @except_wrapper
    def __get_stream(self, client: socket.socket):
        banner = Banner()
        read_banner_len = 0
        banner_length = 24
        read_frame_bytes = 0
        frame_body_len = 0
        data = []
        try:
            while self.capture_stream_flag:
                chunk = client.recv(4096)
                # chunk为空字符则表示minicap服务已经关闭
                if not chunk:
                    break
                cursor = 0
                buf_len = len(chunk)
                while cursor < buf_len:
                    if read_banner_len < banner_length:  # 处理global header, 共24字节, 参考minicap协议
                        map(lambda key, val: banner.__setitem__(key, val), banner.keys, struct.unpack("<2b5ibB", chunk))
                        cursor = buf_len
                        read_banner_len = banner_length
                    elif read_frame_bytes < 4:
                        frame_body_len += (chunk[cursor] << (read_frame_bytes * 8)) >> 0
                        cursor += 1
                        read_frame_bytes += 1
                    else:
                        # 当buf_len - cursor >= frame_body_len时，表示一张图片数据接收完成
                        if buf_len - cursor >= frame_body_len:
                            data.extend(chunk[cursor:cursor + frame_body_len])
                            self.on_image_transferred(data)
                            cursor += frame_body_len
                            frame_body_len = read_frame_bytes = 0
                            data = []
                        else:
                            data.extend(chunk[cursor:buf_len])
                            frame_body_len -= buf_len - cursor
                            read_frame_bytes += buf_len - cursor
                            cursor = buf_len
        except Exception as e:
            Log.error(str(e))
        finally:
            client.close()

    def on_image_transferred(self, data):
        if self.callback:
            self.callback(data)

        if self.save_path:
            filename = datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f') + ".png"
            with open(os.path.join(self.save_path, filename), "wb") as wf:
                for d in data:
                    wf.write(d.to_bytes(1, 'big'))


if __name__ == '__main__':
    cap1 = Minicap(save_path=r"D:\shao\test2")
    cap1.install()
    cap1.start_server()
    cap1.start_capture_stream()
    time.sleep(20)
    cap1.stop_capture_stream()
    cap1.stop_server()

